from flask import Flask, render_template, request, jsonify, session, redirect, url_for, Response
import datetime
import json
import os
import requests
import sqlite3
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 初始化Flask应用
app = Flask(__name__)
app.secret_key = 'your-secret-key-here'  # 用于会话管理的密钥

# 设置后端API地址（指向原FastAPI服务）
BACKEND_API_URL = 'http://localhost:5011'

# 数据库文件路径
DATABASE = 'chat.db'

# 初始化数据库
def init_db():
    with sqlite3.connect(DATABASE) as conn:
        cursor = conn.cursor()
        # 创建topics表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS topics (
                id TEXT PRIMARY KEY,
                title TEXT NOT NULL,
                created_at TEXT NOT NULL,
                user_id TEXT DEFAULT 'default'
            )
        ''')
        # 创建messages表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                topic_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                timestamp TEXT NOT NULL,
                FOREIGN KEY (topic_id) REFERENCES topics (id) ON DELETE CASCADE
            )
        ''')
        conn.commit()

# 从数据库获取所有话题
def get_topics():
    with sqlite3.connect(DATABASE) as conn:
        conn.row_factory = sqlite3.Row  # 返回字典形式的结果
        cursor = conn.cursor()
        cursor.execute('SELECT * FROM topics WHERE user_id = ? ORDER BY created_at DESC', ('default',))
        topics = cursor.fetchall()
        
        # 将结果转换为字典列表，并添加messages字段
        result = []
        for topic in topics:
            topic_dict = dict(topic)
            # 获取该话题的所有消息
            topic_dict['messages'] = get_messages_for_topic(topic_dict['id'])
            result.append(topic_dict)
        
        return result

# 获取特定话题的所有消息
def get_messages_for_topic(topic_id):
    with sqlite3.connect(DATABASE) as conn:
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute('SELECT role, content, timestamp FROM messages WHERE topic_id = ? ORDER BY timestamp ASC', (topic_id,))
        messages = cursor.fetchall()
        return [dict(msg) for msg in messages]

# 保存话题到数据库
def save_topic(topic):
    with sqlite3.connect(DATABASE) as conn:
        cursor = conn.cursor()
        cursor.execute(
            'INSERT OR REPLACE INTO topics (id, title, created_at, user_id) VALUES (?, ?, ?, ?)',
            (topic['id'], topic['title'], topic['created_at'], 'default')
        )
        conn.commit()

# 添加消息到数据库
def add_message(topic_id, message):
    with sqlite3.connect(DATABASE) as conn:
        cursor = conn.cursor()
        cursor.execute(
            'INSERT INTO messages (topic_id, role, content, timestamp) VALUES (?, ?, ?, ?)',
            (topic_id, message['role'], message['content'], message['timestamp'])
        )
        conn.commit()

# 删除话题
def delete_topic(topic_id):
    with sqlite3.connect(DATABASE) as conn:
        cursor = conn.cursor()
        cursor.execute('DELETE FROM topics WHERE id = ?', (topic_id,))
        conn.commit()

# 更新话题标题
def update_topic_title(topic_id, new_title):
    with sqlite3.connect(DATABASE) as conn:
        cursor = conn.cursor()
        cursor.execute('UPDATE topics SET title = ? WHERE id = ?', (new_title, topic_id))
        conn.commit()

# 获取当前话题
def get_current_topic():
    topics = get_topics()
    current_topic_id = session.get('current_topic_id')
    if current_topic_id is not None:
        for topic in topics:
            if topic['id'] == current_topic_id:
                return topic
    return None

# 生成唯一ID
def generate_id():
    return str(datetime.datetime.now().timestamp()).replace('.', '')

# 格式化日期分类
def format_date_category(date_str):
    date = datetime.datetime.strptime(date_str, '%Y-%m-%d %H:%M:%S')
    today = datetime.datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
    tomorrow = today + datetime.timedelta(days=1)
    
    if date >= today:
        return '今天'
    elif date >= today - datetime.timedelta(days=1):
        return '昨天'
    else:
        return '以前'

# 首页路由
@app.route('/')
def index():
    topics = get_topics()
    current_topic = get_current_topic()
    
    # 按日期分类话题
    categorized_topics = {'今天': [], '昨天': [], '以前': []}
    for topic in topics:
        category = format_date_category(topic['created_at'])
        categorized_topics[category].append(topic)
    
    return render_template('index.html', 
                           categorized_topics=categorized_topics,
                           current_topic=current_topic)

# 新建话题路由
@app.route('/new_topic', methods=['POST'])
def new_topic():
    title = request.form.get('title', '未命名话题')
    
    # 创建新话题
    new_topic_obj = {
        'id': generate_id(),
        'title': title,
        'created_at': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }
    
    # 保存到数据库
    save_topic(new_topic_obj)
    
    # 设置当前话题
    session['current_topic_id'] = new_topic_obj['id']
    
    return redirect(url_for('index'))

# 切换话题路由
@app.route('/switch_topic/<topic_id>')
def switch_topic(topic_id):
    # 验证话题是否存在
    topics = get_topics()
    for topic in topics:
        if topic['id'] == topic_id:
            session['current_topic_id'] = topic_id
            break
    return redirect(url_for('index'))

# 删除话题路由
@app.route('/delete_topic/<topic_id>')
def delete_topic_route(topic_id):
    # 删除话题及其消息
    delete_topic(topic_id)
    
    # 如果删除的是当前话题，清除当前话题
    if session.get('current_topic_id') == topic_id:
        session.pop('current_topic_id', None)
        # 如果还有其他话题，切换到第一个话题
        topics = get_topics()
        if topics:
            session['current_topic_id'] = topics[0]['id']
    
    return redirect(url_for('index'))

# 更新话题标题路由
@app.route('/update_topic_title/<topic_id>', methods=['POST'])
def update_topic_title_route(topic_id):
    new_title = request.form.get('title', '未命名话题')
    update_topic_title(topic_id, new_title)
    return redirect(url_for('index'))

# 发送消息路由（AJAX）- 支持流式响应
@app.route('/send_message', methods=['POST'])
def send_message():
    data = request.json
    message = data.get('message')
    
    current_topic = get_current_topic()
    if not current_topic:
        # 如果没有当前话题，创建一个新话题
        current_topic = {
            'id': generate_id(),
            'title': message[:30] if len(message) > 30 else message,  # 使用消息内容作为标题
            'created_at': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        }
        # 保存到数据库
        save_topic(current_topic)
        session['current_topic_id'] = current_topic['id']
    
    # 添加用户消息
    user_message = {
        'role': 'user',
        'content': message,
        'timestamp': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }
    add_message(current_topic['id'], user_message)
    
    def generate():
        try:
            # 准备对话历史用于API调用
            conversation_history = get_messages_for_topic(current_topic['id'])
            
            # 调用后端API获取回复（使用流式请求）
            with requests.post(
                f'{BACKEND_API_URL}/chat',
                json={
                    'message': message,
                    'conversation_history': conversation_history,
                    'stream': True  # 启用流式响应
                },
                stream=True  # 告诉requests库启用流式响应处理
            ) as response:
                if response.status_code == 200:
                    # 保存完整的回复内容
                    full_reply = ""
                    # 处理流式响应中的每一行
                    for line in response.iter_lines():
                        if line:
                            try:
                                # 解析每一行的JSON数据
                                chunk_data = json.loads(line)
                                if 'response' in chunk_data:
                                    full_reply += chunk_data['response']
                                    # 将累积的完整回复通过服务器发送事件(SSE)格式发送给前端
                                    yield f"data: {json.dumps({'type': 'content', 'content': full_reply})}\n\n"
                            except json.JSONDecodeError:
                                # 如果行不是完整的JSON对象，尝试其他处理方式
                                pass
                    
                    # 添加机器人回复到数据库
                    bot_message = {
                        'role': 'assistant',
                        'content': full_reply,
                        'timestamp': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    }
                    add_message(current_topic['id'], bot_message)
                    
                    # 发送结束标志
                    yield f"data: {json.dumps({'type': 'end', 'content': full_reply})}\n\n"
                else:
                    error_msg = f'API调用失败: {response.status_code}'
                    yield f"data: {json.dumps({'type': 'error', 'error': error_msg})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'error': str(e)})}\n\n"
    
    # 返回事件流响应
    return Response(generate(), mimetype='text/event-stream')

# 获取消息历史路由（AJAX）
@app.route('/get_messages/<topic_id>')
def get_messages_route(topic_id):
    messages = get_messages_for_topic(topic_id)
    return jsonify({'messages': messages})

# 保存部分生成内容的路由
@app.route('/save_partial_response', methods=['POST'])
def save_partial_response():
    try:
        data = request.json
        content = data.get('content', '')
        topic_id = data.get('topic_id')
        
        if not topic_id or not content:
            return jsonify({'success': False, 'error': '缺少必要参数'})
        
        # 添加部分机器人回复到数据库
        bot_message = {
            'role': 'assistant',
            'content': content,
            'timestamp': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        }
        add_message(topic_id, bot_message)
        
        return jsonify({'success': True})
    except Exception as e:
        print(f"保存部分回复失败: {str(e)}")
        return jsonify({'success': False, 'error': str(e)})

# 应用启动时初始化数据库
with app.app_context():
    init_db()

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5010)